#!/usr/stsci/pyssgdev/Python-2.5.1/bin/python
import pyfits as p
import numpy as n
import sys

class WCS(object):
    def __init__(self):
        self.crval1 = 0
        self.crval2 = 0
        self.crpix1 = 0
        self.crpix2 = 0
        self.cd1_1 = 0
        self.cd1_2 = 0
        self.cd2_1 = 0
        self.cd2_2 = 0
        self.projection = "TAN"
        self.wcsaxes = 2
        
    def readWCSfile(self,filename):
        file = open(filename)
        lines = file.readlines()
        # Read CRPIX 
        self.crpix1 = float(lines[3].split()[2][:-1])
        self.crpix2 = float(lines[3].split()[3])
        
        # Read CRVAL
        self.crval1 = float(lines[4].split()[2][:-1])
        self.crval2 = float(lines[4].split()[3])
        
        # Read CD Matix
        self.cd1_1 = float(lines[5].split()[2][:-1])
        self.cd1_2 = float(lines[5].split()[3][:-1])
        self.cd2_1 = float(lines[5].split()[4][:-1])
        self.cd2_2 = float(lines[5].split()[5])

        # Read Projection
        proj = lines[2].split()[2].strip()[1:-1]
        if proj == 'TAN':
            self.ctype1 = 'RA---TAN'
            self.ctype2 = 'DEC--TAN'
        else:
            self.ctype1 = proj
            self.ctype2 = proj
        
        file.close()
        
        
    def updateFITSwithWCS(self,inputfilename):
        """
    
        Purpose
        -------
        Update an existing FITS file with valid WCS information.
        
        """
    
        hdulist = p.open(inputfilename,mode="update")
        header = hdulist[0].header
        header.update('WCSAXES',header['NAXIS'])
        header.update('CRPIX1',self.crpix1)
        header.update('CRPIX2',self.crpix2)
        header.update('CRVAL1',self.crval1)
        header.update('CRVAL2',self.crval2)
        header.update('CTYPE1',self.ctype1)
        header.update('CTYPE2',self.ctype2)
        header.update('CD1_1',self.cd1_1)
        header.update('CD1_2',self.cd1_2)
        header.update('CD2_1',self.cd2_1)
        header.update('CD2_2',self.cd2_2)
        hdulist.flush()
        hdulist.close()
        
        
def run(inputfile):
    print "Input: ",inputfile
    file = open(inputfile)
    lines = file.readlines()
    for line in lines:
        print "Processing files: ",line
        newwcs = WCS()
        newwcs.readWCSfile(line.split()[0])
        newwcs.updateFITSwithWCS(line.split()[1])
    file.close()
    
if __name__ == '__main__':
    run(sys.argv[1])

